/* Copyright (c) 2021, National University of Defense Technology. All rights reserved.
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 * 
 *     http://www.apache.org/licenses/LICENSE-2.0
 * 
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

#include <map>
#include <stdlib.h>
#include "csr_transpose.h"
#include "parcsr_transpose.h"

namespace YHAMG
{

#define MPI_TAG 200

void ParCSR_Transpose(const ParCSR_Matrix& A, ParCSR_Matrix& B)
{
	MPI_Comm comm = A.comm;

	int n = A.LocLoc.size[0];
	int m = A.LocLoc.size[1];

	int num_neighbors = A.num_neighbors;
	int* neighbor_ranks = A.neighbor_ranks;
	int* recv_ptr = A.Ext_ptr;
	int* send_ptr = A.send_ptr;
	int* send_list = A.send_list;

	CSR_Matrix A_LocExt_T;
	CSR_Transpose(A.LocExt, A_LocExt_T);

	int* A_LocExt_T_rowptr = A_LocExt_T.rowptr;
	int* A_LocExt_T_colind = A_LocExt_T.colind;
	double* A_LocExt_T_values = A_LocExt_T.values;

	MPI_Request* send_requests = new MPI_Request[num_neighbors];
	MPI_Request* recv_requests = new MPI_Request[num_neighbors];
	MPI_Status status;

	int* send_bufsize = new int[num_neighbors];
	int* recv_bufsize = new int[num_neighbors];

	for (int r = 0; r < num_neighbors; ++r)
		if (send_ptr[r + 1] > send_ptr[r])
			MPI_Irecv(recv_bufsize + r, 1, MPI_INT, neighbor_ranks[r], MPI_TAG, comm, recv_requests + r);

	for (int r = 0; r < num_neighbors; ++r)
	{
		if (recv_ptr[r + 1] > recv_ptr[r])
		{
			int num_rows = recv_ptr[r + 1] - recv_ptr[r];
			int num_nonzeros = A_LocExt_T_rowptr[recv_ptr[r + 1]] - A_LocExt_T_rowptr[recv_ptr[r]];
			send_bufsize[r] = (1 + num_rows + num_nonzeros) * sizeof(int) + num_nonzeros * sizeof(double);
			MPI_Send(send_bufsize + r, 1, MPI_INT, neighbor_ranks[r], MPI_TAG, comm);
		}
	}

	for (int r = 0; r < num_neighbors; ++r)
		if (send_ptr[r + 1] > send_ptr[r])
			MPI_Wait(recv_requests + r, &status);

	char** send_buffer = new char* [num_neighbors];
	char** recv_buffer = new char* [num_neighbors];

	for (int r = 0; r < num_neighbors; ++r)
	{
		if (send_ptr[r + 1] > send_ptr[r])
		{
			recv_buffer[r] = (char*)malloc(recv_bufsize[r]);
			MPI_Irecv(recv_buffer[r], recv_bufsize[r], MPI_PACKED, neighbor_ranks[r], MPI_TAG, comm, recv_requests + r);
		}
	}

	for (int r = 0; r < num_neighbors; ++r)
	{
		if (recv_ptr[r + 1] > recv_ptr[r])
		{
			int num_rows = recv_ptr[r + 1] - recv_ptr[r];
			int starting_row = recv_ptr[r];
			int num_nonzeros = A_LocExt_T_rowptr[recv_ptr[r + 1]] - A_LocExt_T_rowptr[recv_ptr[r]];
			send_buffer[r] = (char*)malloc(send_bufsize[r]);

			int position = 0;
			MPI_Pack(A_LocExt_T_rowptr + starting_row, num_rows + 1, MPI_INT, send_buffer[r], send_bufsize[r], &position, comm);
			MPI_Pack(A_LocExt_T_colind + A_LocExt_T_rowptr[starting_row], num_nonzeros, MPI_INT, send_buffer[r], send_bufsize[r], &position, comm);
			MPI_Pack(A_LocExt_T_values + A_LocExt_T_rowptr[starting_row], num_nonzeros, MPI_DOUBLE, send_buffer[r], send_bufsize[r], &position, comm);

			MPI_Isend(send_buffer[r], send_bufsize[r], MPI_PACKED, neighbor_ranks[r], MPI_TAG, comm, send_requests + r);
		}
	}

	int* B_LocExt_rowptr = new int[m + 1];
	for (int i = 0; i <= m; ++i)
		B_LocExt_rowptr[i] = 0;
	
	int* B_Ext_ptr = new int[num_neighbors + 1];
	B_Ext_ptr[0] = 0;

	CSR_Matrix* A_ExtLoc_T = new CSR_Matrix[num_neighbors];
	std::map<int, int> *map_Ext = new std::map<int, int>[num_neighbors];

	for (int r = 0, cnt = 0; r < num_neighbors; B_Ext_ptr[++r] = cnt)
	{
		if (send_ptr[r + 1] > send_ptr[r])
		{
			MPI_Wait(recv_requests + r, &status);

			int num_rows = send_ptr[r + 1] - send_ptr[r];
			int starting_row = send_ptr[r];
			int num_nonzeros = (recv_bufsize[r] - (num_rows + 1) * sizeof(int)) / (sizeof(int) + sizeof(double));

			int* rowptr = new int[num_rows + 1];
			int* colind = new int[num_nonzeros];
			double* values = new double[num_nonzeros];

			int position = 0;
			MPI_Unpack(recv_buffer[r], recv_bufsize[r], &position, rowptr, num_rows + 1, MPI_INT, comm);
			MPI_Unpack(recv_buffer[r], recv_bufsize[r], &position, colind, num_nonzeros, MPI_INT, comm);
			MPI_Unpack(recv_buffer[r], recv_bufsize[r], &position, values, num_nonzeros, MPI_DOUBLE, comm);

			for (int i = 0, t = rowptr[0]; i <= num_rows; ++i)
				rowptr[i] -= t;

			for (int i = 0; i < num_rows; ++i)
				B_LocExt_rowptr[send_list[starting_row + i]] += rowptr[i + 1] - rowptr[i];

			std::map<int, int> _map_Ext;
			for (int j = 0; j < num_nonzeros; ++j)
			{
				std::pair<std::map<int, int>::iterator, bool> ret = map_Ext[r].insert(std::pair<int, int>(colind[j], cnt + map_Ext[r].size()));
				colind[j] = ret.first->second;
			}

			cnt += map_Ext[r].size();

			A_ExtLoc_T[r].size[0] = num_rows;
			A_ExtLoc_T[r].size[1] = map_Ext[r].size();
			A_ExtLoc_T[r].rowptr = rowptr;
			A_ExtLoc_T[r].colind = colind;
			A_ExtLoc_T[r].values = values;

			free(recv_buffer[r]);
		}
	}

	for (int r = 0; r < num_neighbors; ++r)
	{
		if (recv_ptr[r + 1] > recv_ptr[r])
		{
			MPI_Wait(send_requests + r, &status);
			free(send_buffer[r]);
		}
	}

	delete[] recv_bufsize;
	delete[] send_bufsize;
	delete[] recv_buffer;
	delete[] recv_requests;
	delete[] send_buffer;
	delete[] send_requests;

	int* B_Ext_map = new int[B_Ext_ptr[num_neighbors]];
	for (int r = 0; r < num_neighbors; ++r)
	{
		if (send_ptr[r + 1] > send_ptr[r])
		{
			for (std::map<int, int>::iterator iter = map_Ext[r].begin(); iter != map_Ext[r].end(); ++iter)
				B_Ext_map[iter->second] = iter->first;
		}
	}

	for (int i = 0; i < m; ++i)
		B_LocExt_rowptr[i + 1] += B_LocExt_rowptr[i];

	int* B_LocExt_colind = new int[B_LocExt_rowptr[m]];
	double* B_LocExt_values = new double[B_LocExt_rowptr[m]];

	for (int r = num_neighbors - 1; r >= 0; --r)
	{
		if (send_ptr[r + 1] > send_ptr[r])
		{
			int starting_row = send_ptr[r];
			int num_rows = A_ExtLoc_T[r].size[0];
			int* rowptr = A_ExtLoc_T[r].rowptr;
			int* colind = A_ExtLoc_T[r].colind;
			double* values = A_ExtLoc_T[r].values;

			for (int i = num_rows - 1; i >= 0; --i)
			{
				int k = B_LocExt_rowptr[send_list[starting_row + i]];
				for (int j = rowptr[i + 1] - 1; j >= rowptr[i]; --j)
				{
					B_LocExt_colind[--k] = colind[j];
					B_LocExt_values[k] = values[j];
				}
				B_LocExt_rowptr[send_list[starting_row + i]] = k;
			}
		}
	}

	delete[] A_ExtLoc_T;

	int* B_neighbors = new int[num_neighbors];
	for (int r = 0; r < num_neighbors; ++r)
		B_neighbors[r] = neighbor_ranks[r];

	CSR_Matrix B_LocLoc;
	CSR_Transpose(A.LocLoc, B_LocLoc);
	
	B_LocLoc.ref = 1;

	B.Free();

	B.comm = comm;
	B.LocLoc.size[0] = B_LocLoc.size[0];
	B.LocLoc.size[1] = B_LocLoc.size[1];
	B.LocLoc.rowptr = B_LocLoc.rowptr;
	B.LocLoc.colind = B_LocLoc.colind;
	B.LocLoc.values = B_LocLoc.values;
	B.LocExt.size[0] = m;
	B.LocExt.size[1] = B_Ext_ptr[num_neighbors];
	B.LocExt.rowptr = B_LocExt_rowptr;
	B.LocExt.colind = B_LocExt_colind;
	B.LocExt.values = B_LocExt_values;
	B.num_neighbors = num_neighbors;
	B.neighbor_ranks = B_neighbors;
	B.Ext_ptr = B_Ext_ptr;
	B.Ext_map = B_Ext_map;
}

}